import numpy as np
from langchain_core.prompts import PromptTemplate

# 示例
examples = [
    {"question": "基金今天适合买入吗？", "answer": "可以适量买入"},
    {"question": "基金的风险是什么？", "answer": "市场波动风险"},
    {"question": "基金经理是谁？", "answer": "张三"},
]

example_prompt = PromptTemplate(
    input_variables=["question", "answer"],
    template="Q: {question}\nA: {answer}"
)

# 本地 embedding
def local_embed(text):
    return np.array([len(text)])  # 简单示例

# 生成示例向量
example_vectors = [local_embed(ex["question"]) for ex in examples]

# 自定义 selector
class LocalSemanticSelector:
    def __init__(self, examples, vectors, k=2):
        self.examples = examples
        self.vectors = vectors
        self.k = k

    def select_examples(self, input_str):
        input_vec = local_embed(input_str)
        similarities = [
            np.dot(input_vec, v) / (np.linalg.norm(input_vec)*np.linalg.norm(v))
            for v in self.vectors
        ]
        topk_idx = np.argsort(similarities)[-self.k:][::-1]
        return [self.examples[i] for i in topk_idx]

# 使用
selector = LocalSemanticSelector(examples, example_vectors, k=2)
selected = selector.select_examples("基金风险有哪些？")
for s in selected:
    print(s)
